import json
import os

from memory_systems.agent_stats_db import AgentStatsDB

class AgentInfoDB:
    def __init__(self, db_path=None):
        if db_path is None:
            # Assuming BASE_DIR is the project root
            BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
            self.db_path = os.path.join(BASE_DIR, "agents", "agent_metadata.json")
            self.stats_db_path = os.path.join(BASE_DIR, "memory_systems", "agent_stats.db")
        else:
            self.db_path = db_path
            # If a custom db_path is provided, assume stats_db is in the same directory or handle as needed
            self.stats_db_path = os.path.join(os.path.dirname(db_path), "agent_stats.db")
        self._initialize_db()
        self.agent_stats_db = AgentStatsDB(db_path=self.stats_db_path)

    def _initialize_db(self):
        if not os.path.exists(os.path.dirname(self.db_path)):
            os.makedirs(os.path.dirname(self.db_path))
        if not os.path.exists(self.db_path):
            with open(self.db_path, 'w', encoding='utf-8') as f:
                json.dump({}, f, indent=4, ensure_ascii=False)

    def _load_metadata(self):
        try:
            with open(self.db_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            print(f"Error decoding JSON from {self.db_path}. Returning empty metadata.")
            return {}
        except Exception as e:
            print(f"Error loading agent metadata from {self.db_path}: {e}. Returning empty metadata.")
            return {}

    def _save_metadata(self, metadata):
        try:
            with open(self.db_path, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, indent=4, ensure_ascii=False)
        except Exception as e:
            print(f"Error saving agent metadata to {self.db_path}: {e}")

    def add_agent(self, agent_name: str, description: str, version: str = "1.0", enabled: bool = True):
        metadata = self._load_metadata()
        if agent_name not in metadata:
            metadata[agent_name] = {
                "name": agent_name,
                "description": description,
                "version": version,
                "enabled": enabled
            }
            self._save_metadata(metadata)
            # Initialize stats in the separate DB
            self.agent_stats_db.increment_run_count(agent_name, success=False) # Initialize with 0 runs, 0 success
            return True
        return False # Agent already exists

    def update_agent_info(self, agent_name: str, **kwargs):
        metadata = self._load_metadata()
        if agent_name in metadata:
            for key, value in kwargs.items():
                if key in metadata[agent_name] and key not in ["run_count", "success_count"]: # Only update existing keys, exclude stats
                    metadata[agent_name][key] = value
            self._save_metadata(metadata)
            return True
        return False # Agent not found

    def get_agent_info(self, agent_name: str):
        metadata = self._load_metadata()
        info = metadata.get(agent_name)
        if info:
            stats = self.agent_stats_db.get_agent_stats(agent_name)
            if stats:
                info.update(stats) # Merge stats into info
        return info

    def get_agent_stats(self, agent_name: str):
        return self.agent_stats_db.get_agent_stats(agent_name)

    def get_all_agents(self, include_disabled: bool = False):
        metadata = self._load_metadata()
        all_stats = self.agent_stats_db.get_all_agent_stats()
        stats_map = {stat["agent_name"]: stat for stat in all_stats}

        combined_agents = {}
        for agent_name, info in metadata.items():
            # Only include agents that are explicitly enabled or if the 'enabled' key is missing (default to enabled)
            # Or if include_disabled is True
            if include_disabled or info.get("enabled", True):
                combined_info = info.copy()
                if agent_name in stats_map:
                    combined_info.update(stats_map[agent_name])
                combined_agents[agent_name] = combined_info
        return combined_agents
        return combined_agents

    def increment_run_count(self, agent_name: str, success: bool = True):
        return self.agent_stats_db.increment_run_count(agent_name, success)

    def enable_agent(self, agent_name: str):
        return self.update_agent_info(agent_name, enabled=True)

    def disable_agent(self, agent_name: str):
        if agent_name.lower() == "rootagent":
            print("Cannot disable rootagent.")
            return False
        return self.update_agent_info(agent_name, enabled=False)

    def delete_agent(self, agent_name: str):
        metadata = self._load_metadata()
        if agent_name in metadata:
            del metadata[agent_name]
            self._save_metadata(metadata)
            # Also delete stats from the separate DB
            self.agent_stats_db.delete_agent_stats(agent_name)
            return True
        return False
